package com.TanNgee.netty.kyro.serialize;

/**
 * @author 陈毅
 * @create 2022/3/11 0:20
 */

import com.TanNgee.netty.kyro.dto.RpcRequest;
import com.TanNgee.netty.kyro.dto.RpcResponse;
import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;

/**
 * 自定义kryo序列化接口实现类
 */
public class KryoSerializer implements Serializer {
    //    由于Kryo不是线程安全，所以每个线程都要有自己的Kryo，Input和Output实例
    // 使用ThreadLocal 来存放
    private static final ThreadLocal<Kryo> kryoThreadLocal = ThreadLocal.withInitial(() -> {
        Kryo kryo = new Kryo();

        kryo.register(RpcResponse.class);
        kryo.register(RpcRequest.class);

        kryo.setReferences(true);//默认值为true,是否关闭注册行为,关闭之后可能存在序列化问题，一般推荐设置为 true
        kryo.setRegistrationRequired(false);//默认值为false,是否关闭循环引用，可以提高性能，但是一般不推荐设置为 true
        return kryo;
    });


    /**
     * 序列化
     *
     * @param obj 要序列化的对象
     * @return
     */
    @Override
    public byte[] serialize(Object obj) {
        try (ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
             Output output = new Output(byteArrayOutputStream)
        ) {
//            获取kryo对象
            Kryo kryo = kryoThreadLocal.get();
            // 将Object对象转换为byte数组
            kryo.writeObject(output, obj);

            kryoThreadLocal.remove();

            return output.toBytes();
        } catch (IOException e) {
            // 抛出自定义异常
            throw new SerializeException("序列化失败");
        }
    }


    /**
     * 反序列化
     *
     * @param bytes 被序列化的字节数组
     * @param clazz 要反序列化的类
     * @param <T>
     * @return
     */
    @Override
    public <T> T deserialize(byte[] bytes, Class<T> clazz) {

        try (ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes);
             Input input = new Input(byteArrayInputStream)) {
            // 获取Kryo对象
            Kryo kryo = kryoThreadLocal.get();

            // 反序列化
            Object o = kryo.readObject(input, clazz);

            kryoThreadLocal.remove();
            return clazz.cast(o);

        } catch (IOException e) {
            // 抛出自定义异常
            throw new SerializeException("反序列化失败");
        }
    }
}
